#include "file_crypto.h"
#include "sm3.h"
#include "sm4.h"
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <sys/stat.h>

int genKey(const char *cpuId, int cpuIdSize, const char *gpuId, int gpuIdSize, char *key)
{
	unsigned char cpuKey[32 +1] = {0};
	unsigned char gpuKey[32 +1] = {0};
	int n;
	
	if((cpuId == NULL) || (gpuId == NULL) || (key == NULL))
		return R_NULLPOINTER;

	if((cpuIdSize <= 0) || (gpuIdSize <= 0))
		return R_SIZEERR;

	sm3((unsigned char *)cpuId, cpuIdSize, cpuKey);
	sm3((unsigned char *)gpuId, gpuIdSize, gpuKey);

	for(n = 0; n < 16; n++)
		key[n] = cpuKey[n] ^ gpuKey[n + 16];

	return R_OK;
}

#define SM4_ENCRYPT     1
#define SM4_DECRYPT     0
#define PADDING_SIZE 16
#define FILE_BUF_SIZE 1024 * 1024 * 10
int symFileEnc(const char *key, const char *plainFileName, const char *cipherFileName)
{
	FILE *inFp = NULL;
	FILE *outFp = NULL;
	unsigned char iv[16 + 1] = {0};
	char *in = NULL;
	char *out = NULL;
	unsigned long fileSize;
	int nret = -1;
	struct stat st;
	sm4_context ctx;

	if((key == NULL) || (plainFileName == NULL) || (cipherFileName == NULL))
	{
		nret = R_INARGERR;
		goto err;
	}

	memset(iv, 0x00, sizeof(iv));

	stat(plainFileName, &st);
	fileSize = st.st_size;

	inFp = fopen(plainFileName, "rb");
	if(!inFp)
	{
		nret = R_OPEN_FILE_ERR;
		goto err;
	}

	remove(cipherFileName);
	outFp = fopen(cipherFileName, "wb");
	if(!outFp)
	{
		nret = R_OPEN_FILE_ERR;
		goto err;
	}


	if((in = (char *)malloc(FILE_BUF_SIZE + PADDING_SIZE)) == NULL)
	{
		nret = R_MEMORYERR;
		goto err;
	}

	if((out = (char *)malloc(FILE_BUF_SIZE + PADDING_SIZE + 1)) == NULL)
	{
		nret = R_MEMORYERR;
		goto err;
	}

	sm4_setkey_enc(&ctx, (unsigned char *)key);

	while(!feof(inFp))
	{
		int padding = 0, len = 0, nread = 0;
		int isPadding = 0;

		if(fileSize > FILE_BUF_SIZE)
			nread = FILE_BUF_SIZE;
		else
		{
			nread = fileSize;
			isPadding = 1;
		}

		memset(in, 0x00, FILE_BUF_SIZE + PADDING_SIZE);
		memset(out, 0x00, FILE_BUF_SIZE + PADDING_SIZE);
		if(fread(in, nread, 1, inFp) != 1)
		{
			nret = R_READ_FILE_ERR;
			goto err;
		}

		len = nread;

		if(isPadding == 1)
		{
			padding = 16 - len % 16;
			memset(in + len, padding, padding);
			len += padding;
			in[len] = '\0';
		}

		sm4_crypt_cbc(&ctx, SM4_ENCRYPT, len, iv, (unsigned char *)in, (unsigned char *)out);

		if(fwrite(out, len, 1, outFp) != 1)
		{
			nret = R_WRITE_FILE_ERR;
			goto err;
		}

		fileSize -= nread;
		if(fileSize <= 0)
			break;
	}

	nret = R_OK;
err:
	if(inFp)
		fclose(inFp);

	if(outFp)
		fclose(outFp);

	if(in)
		free(in);

	if(out)
		free(out);

	return nret;
}

int symFileDec(const char *key, const char *cipherFileName, const char *plainFileName)
{
	FILE *inFp = NULL;
	FILE *outFp = NULL;
	char *in = NULL;
	char *out = NULL;
	int nret = -1;
	unsigned char iv[16 + 1] = {0};
	struct stat st;
	unsigned long fileSize;
	sm4_context ctx;

	if((key == NULL) || (plainFileName == NULL) || (cipherFileName == NULL))
	{
		nret = R_INARGERR;
		goto err;
	}

	memset(iv, 0x00, sizeof(iv));

	stat(cipherFileName, &st);
	fileSize = (st.st_size - 16);

	inFp = fopen(cipherFileName, "rb");
	if(!inFp)
	{
		nret = R_OPEN_FILE_ERR;
		goto err;
	}

	remove(plainFileName);
	outFp = fopen(plainFileName, "wb");
	if(!outFp)
	{
		nret = R_OPEN_FILE_ERR;
		goto err;
	}

	if((in = (char *)malloc(FILE_BUF_SIZE + PADDING_SIZE)) == NULL)
	{
		nret = R_MEMORYERR;
		goto err;
	}

	if((out = (char *)malloc(FILE_BUF_SIZE + PADDING_SIZE)) == NULL)
	{
		nret = R_MEMORYERR;
		goto err;
	}

	sm4_setkey_dec(&ctx, (unsigned char *)key);
	
	while(!feof(inFp))
	{
		int padding, len, nread;
		int isPadding = 0;

		if(fileSize > FILE_BUF_SIZE)
			nread = FILE_BUF_SIZE;
		else
		{
			nread = fileSize;
			isPadding = 1;
		}

		memset(in, 0x00, FILE_BUF_SIZE);
		memset(out, 0x00, FILE_BUF_SIZE + PADDING_SIZE);
		if(fread(in, nread, 1, inFp) != 1)
		{

			nret = R_READ_FILE_ERR;
			goto err;
		}

		len = nread;

		sm4_crypt_cbc(&ctx, SM4_DECRYPT, len, iv, (unsigned char *)in, (unsigned char *)out);

		if(isPadding == 1)
		{
			padding = out[len - 1];
			out[len - padding] = '\0';
			len -= padding;
		}

		if(fwrite(out, len, 1, outFp) != 1)
		{
			nret = R_WRITE_FILE_ERR;
			goto err;
		}

		fileSize -= nread;
		if(fileSize <= 0)
			break;
	}

	nret = R_OK;
err:
	if(inFp)
		fclose(inFp);

	if(outFp)
		fclose(outFp);

	if(in)
		free(in);

	if(out)
		free(out);

	return nret;
}
